from deep_translator import GoogleTranslator
import os
from pathlib import Path
import shutil
import pypdf
import re
from io import BytesIO

def translate_pdf_file(input_path, output_path=None, source_lang='auto', target_lang='zh-CN', from_page=1, to_page=None):
    """
    翻译PDF文件（支持多页，支持页码范围选择，保持PDF格式）
    
    参数:
        input_path (str): 输入PDF文件路径
        output_path (str, optional): 输出文件路径，默认为None（自动生成）
        source_lang (str): 源语言代码，默认为auto自动检测
        target_lang (str): 目标语言代码，默认为zh-CN
        from_page (int): 起始页码（从1开始），默认为1
        to_page (int, optional): 结束页码（包含），默认为None（表示最后一页）
        
    返回:
        tuple: (翻译后的文本, 输出文件路径)
    """
    try:
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 检查页码范围
        if from_page < 1:
            raise ValueError(f"起始页码必须大于等于1，当前值: {from_page}")
        
        # 如果未指定输出路径，则自动生成
        if output_path is None:
            input_file = Path(input_path)
            filename = input_file.stem
            page_range = ""
            if from_page > 1 or to_page is not None:
                page_range = f"_p{from_page}"
                if to_page is not None:
                    page_range += f"-{to_page}"
            output_path = os.path.join("output_files", f"{filename}{page_range}_{target_lang}.pdf")
        elif not output_path.lower().endswith('.pdf'):
            output_path += '.pdf'
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 创建翻译器
        translator = GoogleTranslator(source=source_lang, target=target_lang)
        
        # 准备文本输出路径以保存所有翻译文本
        txt_output_path = output_path.replace('.pdf', '.txt')
        all_translated_text = ""
        
        # 打开PDF文件
        with open(input_path, 'rb') as pdf_file:
            pdf_reader = pypdf.PdfReader(pdf_file)
            total_pages = len(pdf_reader.pages)
            
            # 处理页码范围
            if to_page is None or to_page > total_pages:
                to_page = total_pages
            
            if from_page > total_pages:
                raise ValueError(f"起始页码超出PDF总页数。PDF共有 {total_pages} 页，指定的起始页码为 {from_page}")
            
            print(f"正在处理PDF文件，共 {total_pages} 页，翻译范围: {from_page} - {to_page} 页...")
            
            # 简化处理方式：首先复制原始PDF到输出路径
            shutil.copy2(input_path, output_path)
            
            # 逐页提取和翻译文本（仅用于生成文本文件和返回翻译文本）
            for page_num in range(from_page - 1, to_page):
                # 获取当前页
                page = pdf_reader.pages[page_num]
                # 提取文本
                page_text = page.extract_text()
                
                if page_text.strip():  # 确保页面有内容
                    try:
                        # 分段翻译，处理长文本
                        translated_page = translate_long_text(page_text, translator)
                        all_translated_text += f"\n\n--- 第 {page_num + 1} 页 ---\n\n"
                        all_translated_text += translated_page
                        print(f"已完成第 {page_num + 1} 页翻译")
                    except Exception as e:
                        all_translated_text += f"\n\n--- 第 {page_num + 1} 页（翻译失败）---\n\n"
                        all_translated_text += f"翻译错误: {str(e)}\n原文: {page_text[:100]}..."
                        print(f"第 {page_num + 1} 页翻译失败: {str(e)}")
        
        # 将翻译结果保存到文本文件
        with open(txt_output_path, 'w', encoding='utf-8') as f:
            f.write(all_translated_text)
        
        print(f"PDF文件翻译完成，原始PDF文件已保存到: {output_path}")
        print(f"翻译文本已保存到: {txt_output_path}")
        return all_translated_text, output_path
    
    except Exception as e:
        raise Exception(f"PDF文件翻译失败: {str(e)}")

def translate_long_text(text, translator, max_chunk_size=4000):
    """
    分段翻译长文本
    
    参数:
        text (str): 要翻译的文本
        translator: GoogleTranslator实例
        max_chunk_size (int): 每段最大字符数
    
    返回:
        str: 翻译后的完整文本
    """
    # 如果文本较短，直接翻译
    if len(text) <= max_chunk_size:
        return translator.translate(text)
    
    # 分段处理长文本
    chunks = []
    current_chunk = ""
    
    # 按句子拆分文本
    sentences = re.split(r'(?<=[.!?])\s+', text)
    
    for sentence in sentences:
        # 如果当前块加上新句子不超过最大大小，添加到当前块
        if len(current_chunk) + len(sentence) <= max_chunk_size:
            current_chunk += sentence + " "
        else:
            # 否则，将当前块添加到chunks，并开始新块
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = sentence + " "
    
    # 添加最后一个块
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    # 翻译每个块并合并
    translated_chunks = []
    for i, chunk in enumerate(chunks):
        translated_chunk = translator.translate(chunk)
        translated_chunks.append(translated_chunk)
    
    # 合并所有翻译后的块
    return " ".join(translated_chunks)

def translate_docx_file(input_path, output_path=None, source_lang='auto', target_lang='zh-CN'):
    """
    翻译DOCX文件（保持DOCX格式）
    
    参数:
        input_path (str): 输入DOCX文件路径
        output_path (str, optional): 输出文件路径，默认为None（自动生成）
        source_lang (str): 源语言代码，默认为auto自动检测
        target_lang (str): 目标语言代码，默认为zh-CN
        
    返回:
        tuple: (翻译后的文本, 输出文件路径)
    """
    try:
        # 确保输入文件存在
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"输入文件不存在: {input_path}")
        
        # 如果未指定输出路径，则自动生成
        if output_path is None:
            input_file = Path(input_path)
            filename = input_file.stem
            output_path = os.path.join("output_files", f"{filename}_{target_lang}.docx")
        elif not output_path.lower().endswith('.docx'):
            output_path += '.docx'
        
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 先复制原始文件到输出路径，以防万一文本处理失败
        shutil.copy2(input_path, output_path)
        
        # 创建翻译器
        translator = GoogleTranslator(source=source_lang, target=target_lang)
        
        # 准备文本输出路径
        txt_output_path = output_path.replace('.docx', '.txt')
        all_translated_text = ""
        
        # 使用python-docx库处理文档
        try:
            import docx
            
            # 打开复制后的文档
            doc = docx.Document(output_path)
            
            # 翻译段落
            for para in doc.paragraphs:
                if para.text.strip():
                    original_text = para.text
                    try:
                        translated_text = translate_long_text(original_text, translator)
                        all_translated_text += translated_text + "\n\n"
                    except Exception as e:
                        print(f"警告: 段落翻译失败: {str(e)[:100]}...")
                        # 保留原始文本
                        all_translated_text += f"[翻译失败] {original_text}\n\n"
            
            # 翻译表格内容（仅文本记录，不修改原文档）
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        if cell.text.strip():
                            original_text = cell.text
                            try:
                                translated_text = translate_long_text(original_text, translator)
                                all_translated_text += translated_text + " | "
                            except Exception as e:
                                print(f"警告: 表格单元格翻译失败: {str(e)[:100]}...")
                                # 保留原始文本
                                all_translated_text += f"[翻译失败] {original_text} | "
                    all_translated_text += "\n"
                all_translated_text += "\n\n"
            
            # 保存文本版本
            with open(txt_output_path, 'w', encoding='utf-8') as f:
                f.write(all_translated_text)
            
            print(f"DOCX文件翻译完成，原始DOCX文件已保存到: {output_path}")
            print(f"翻译文本已保存到: {txt_output_path}")
            
            return all_translated_text, output_path
            
        except ImportError:
            # 如果未安装python-docx，则无法处理DOCX
            print("错误: 未安装python-docx库，无法处理DOCX文件。")
            print("请安装: pip install python-docx")
            
            # 尝试生成文本版本
            try:
                with open(input_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
                    translated_text = translate_long_text(text, translator)
                    with open(txt_output_path, 'w', encoding='utf-8') as out:
                        out.write(translated_text)
                    return translated_text, output_path
            except Exception:
                pass
            
            raise ImportError("需要python-docx库来处理DOCX文件。请安装: pip install python-docx")
        
    except Exception as e:
        raise Exception(f"DOCX文件翻译失败: {str(e)}") 